import network
import socket
import machine
import time
from machine import I2S, Pin, I2C
import struct
import array

# 如果没有 ssd1306 库，请确保已上传 ssd1306.py 到 Pico
try:
    import ssd1306
    import font
    OLED_AVAILABLE = True
except ImportError:
    OLED_AVAILABLE = False
    print("OLED libraries not found. Running without display.")

# ================= 配置区域 (CONFIGURATION) =================
# 1. WiFi Settings
SSID = "Astro_RDX"       # <--- 修改这里
PASSWORD = "17060840"    # <--- 修改这里

# 2. PC Target Settings (Echo Server IP)
PC_IP = "192.168.43.151" # <--- 修改为你的电脑 IP
PORT = 50005

# 3. Pin Definitions (Raspberry Pi Pico W)
PIN_WS = 14      # LRCL / WS
PIN_SCK = 13     # BCLK / SCK
PIN_SD_MIC = 12  # Mic DOUT (SD)
PIN_SD_SPK = 15  # Speaker DIN (SD)
PIN_BTN = 11     # Button
PIN_SDA = 16     # OLED SDA
PIN_SCL = 17     # OLED SCL

# 4. Audio Settings
# ICS-43434 和 MAX98357A 均支持 32-bit slot
SAMPLE_RATE = 16000
SAMPLE_BITS = 32 # 关键修改：统一使用 32 位宽
CHUNK_SIZE = 1024 # 字节数 (256个采样点 * 4字节)
# ==========================================================

# --- 1. OLED Initialization ---
oled = None
if OLED_AVAILABLE:
    try:
        print("[INIT] Initializing OLED...")
        i2c = I2C(0, scl=Pin(PIN_SCL), sda=Pin(PIN_SDA), freq=400000)
        oled = ssd1306.SSD1306_I2C(128, 64, i2c)
    except Exception as e:
        print(f"OLED Init Failed: {e}")

def log(msg, line2=""):
    """ 同时输出到 串口 和 OLED """
    print(f"[{time.ticks_ms()/1000:.1f}s] {msg} {line2}")
    if oled:
        try:
            oled.fill(0)
            # 简单的文本显示，根据你的 font 库调整
            if 'font' in globals():
                font.draw(oled, msg, 0, 0)
                if line2: font.draw(oled, line2, 0, 12)
            else:
                oled.text(msg, 0, 0)
                oled.text(line2, 0, 12)
            oled.show()
        except: pass

# --- 2. 核心修复：正确的音量调节 ---
def adjust_volume(data, shift=4):
    """
    使用 array 模块代替 memoryview.cast 来处理 32-bit 数据
    """
    if shift == 0:
        return data
        
    # 1. 将字节串转换为 32位有符号整数数组 ('i')
    # MicroPython 的 array 初始化会复制数据，虽然不如 memoryview 高效，
    # 但兼容性最好，且对于 1024 字节的数据量，速度足够快。
    samples = array.array('i', data)
    
    # 2. 遍历数组进行位移操作
    length = len(samples)
    for i in range(length):
        # 算术右移，保留符号位
        samples[i] = samples[i] >> shift
        
    # 3. 将处理后的数组转回字节串
    return bytes(samples)

log("System Booting...")

# --- 3. Hardware Initialization ---
led = machine.Pin("LED", machine.Pin.OUT)
btn = machine.Pin(PIN_BTN, machine.Pin.IN, machine.Pin.PULL_UP)

# I2S Initialization
# 关键修复：IN 和 OUT 全部统一为 32 bits
# 这样 Mic 的 24-bit 数据会自动对齐到高位
# Speaker 也会自动读取 32-bit 数据的高 16/24 位进行播放
print("Configuring I2S (32-bit)...")

audio_in = I2S(0,
    sck=Pin(PIN_SCK), ws=Pin(PIN_WS), sd=Pin(PIN_SD_MIC),
    mode=I2S.RX, 
    bits=SAMPLE_BITS, # 32
    format=I2S.MONO, 
    rate=SAMPLE_RATE, 
    ibuf=CHUNK_SIZE * 4 
)

audio_out = I2S(1,
    sck=Pin(PIN_SCK), ws=Pin(PIN_WS), sd=Pin(PIN_SD_SPK),
    mode=I2S.TX, 
    bits=SAMPLE_BITS, # 32 (匹配 Mic，也匹配 MAX98357A 的灵活性)
    format=I2S.MONO, 
    rate=SAMPLE_RATE, 
    ibuf=CHUNK_SIZE * 4
)

# --- 4. WiFi Connection ---
log("Connecting WiFi...", SSID)
wlan = network.WLAN(network.STA_IF)
wlan.active(True)
# 禁用 WiFi 节能模式，降低延迟 (Raspberry Pi Pico W 特定)
wlan.config(pm=0xa11140) 
wlan.connect(SSID, PASSWORD)

retry = 0
while not wlan.isconnected():
    time.sleep(1)
    retry += 1
    print(".", end="")
    if retry > 20:
        log("WiFi Failed!", "Check Config")
        # 即使 WiFi 失败，不要死循环，方便调试
        break

print("\nIP:", wlan.ifconfig()[0])
log("WiFi OK", wlan.ifconfig()[0])
led.on()
time.sleep(1)

# --- 5. UDP Connection ---
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setblocking(False) 
server_addr = (PC_IP, PORT)

mic_buffer = bytearray(CHUNK_SIZE)

log("Ready", "Hold BTN to Talk")

# --- 6. Main Loop ---
try:
    while True:
        # ====================
        # A. Recording Logic (Send)
        # ====================
        if btn.value() == 0:
            log("Recording...", "Speak Now!")
            led.off()
            
            # Send Start Signal
            try:
                sock.sendto(b'START', server_addr)
            except OSError: pass
            
            packet_count = 0
            
            # 清空之前的接收缓冲，防止听到旧声音
            try:
                while True:
                    sock.recv(1024)
            except OSError: pass

            while btn.value() == 0:
                # 读取 32位 原始数据
                num_read = audio_in.readinto(mic_buffer)
                
                if num_read > 0:
                    # 直接发送原始 32位 数据到 PC
                    # PC 端如果只是 Echo，不需要关心具体格式，原样发回即可
                    try:
                        sent = sock.sendto(mic_buffer[:num_read], server_addr)
                        packet_count += 1
                    except OSError as e:
                        # WiFi 拥塞丢包是正常的，忽略
                        pass
                    
                    if packet_count % 50 == 0:
                        print(".", end="")
            
            print(f"\nSent {packet_count} packets")
            try:
                sock.sendto(b'STOP', server_addr)
            except: pass
            
            log("Sent!", "Waiting Echo...")
            led.on()
            time.sleep(0.1) # Debounce
        
        # ====================
        # B. Playback Logic (Receive)
        # ====================
        try:
            # 接收回声数据
            data, addr = sock.recvfrom(CHUNK_SIZE * 2)
            
            if data:
                if data == b'END_OF_STREAM':
                    log("Finished", "Hold BTN")
                elif len(data) > 0:
                    # 【重要】在这里进行音量衰减
                    # 因为 Mic 输入的是全幅度的 32位 信号，MAX98357A 增益很高
                    # 建议 shift=3 (/8) 或 shift=4 (/16) 开始测试
                    quiet_data = adjust_volume(data, shift=3)
                    
                    # 写入 I2S (32-bit 模式)
                    audio_out.write(quiet_data)
                    
        except OSError:
            # Non-blocking 模式下没有数据会抛出异常，忽略即可
            pass
        except Exception as e:
            print("RX Error:", e)

except KeyboardInterrupt:
    log("Stopped", "By User")
    audio_in.deinit()
    audio_out.deinit()
    print("I2S Deinitialized")